""" Taken and adapted from https://github.com/cyclomon/3dbraingen """
import glob
import os
import numpy as np
import nibabel as nib
import torch
import torchio as tio
from torch.utils.data import Dataset
from skimage.transform import resize

class ADNIDataset(Dataset):
    # 修改 __init__ 方法
    def __init__(self, file_list, augmentation=False):
        self.file_names = file_list  # 不再扫描目录，直接使用传入的列表
        self.augmentation = augmentation

    def __len__(self):
        return len(self.file_names)

    def roi_crop(self, image):
        # ... (这部分代码保持不变)
        if image.ndim == 4:
            # 只有当数组是4维时，才执行降维操作
            image = image[:, :, :, 0]
        mask = image > 0
        coords = np.argwhere(mask)
        x0, y0, z0 = coords.min(axis=0)
        x1, y1, z1 = coords.max(axis=0) + 1
        cropped = image[x0:x1, y0:y1, z0:z1]
        padded_crop = tio.CropOrPad(
            np.max(cropped.shape))(cropped.copy()[None])
        padded_crop = np.transpose(padded_crop, (1, 2, 3, 0))
        return padded_crop

    def __getitem__(self, index):
        # ... (这部分代码保持不变)
        path = self.file_names[index]

        img_original = nib.load(path)
        # 2. 使用 as_closest_canonical() 自动重采样到标准方向 (RAS)
        #    这会返回一个新的nii对象，其数据方向是标准化的
        img_canonical = nib.as_closest_canonical(img_original)

        # 3. 从这个标准化的对象中获取数据
        #    使用 .get_fdata() 可以确保字节序和数据类型正确
        img_data = img_canonical.get_fdata(dtype=np.float32)

        # img = nib.load(path)
        # img_data = img.get_fdata(dtype=np.float32)
        # img = np.swapaxes(img_data, 1, 2)
        # img = np.swapaxes(img.get_data(), 1, 2)
        # img = np.flip(img, 1)
        # img = np.flip(img, 2)

        img = self.roi_crop(image=img_data)
        sp_size = 64
        img = resize(img, (sp_size, sp_size, sp_size), mode='constant')

        # 防止出现特小特大值,归一化为[-1,1]
        min_val = np.percentile(img, 0.5)
        max_val = np.percentile(img, 99.5)

        # 防止除以零
        if max_val - min_val > 1e-6:
            # 2. 将数据裁剪到这个百分位范围内，去除极端异常值
            img = np.clip(img, min_val, max_val)

            # 3. 线性缩放到 [0, 1] 范围
            img = (img - min_val) / (max_val - min_val)
        else:
            # 如果图像是全黑的或强度非常单一，直接设为0
            img.fill(0)

        if self.augmentation:
            random_n = torch.rand(1)
            random_i = 0.3*torch.rand(1)[0]+0.7
            if random_n[0] > 0.5:
                img = np.flip(img, 0)
            img = img*random_i.data.cpu().numpy()

        imageout = torch.from_numpy(img.copy()).float().view(1, sp_size, sp_size, sp_size)
        imageout = imageout*2-1
        return {'data': imageout}